#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>

#include <fstream>
#include <random>

#include <bthread/bthread.h>
#include <gtest/gtest.h>

#include "util/memory/array/simple_recycle_array.h"

DEFINE_int32(simple_recycle_array_test_tasks, 128, "num of threads or bthreads");
DEFINE_int32(simple_recycle_array_test_data_size, 1024000, "data file size");
DEFINE_int32(simple_recycle_array_test_doc_num, 102400, "doc num");

void InitOffsetFile(const std::string& offset_filename) {
    std::ofstream ofs(offset_filename.data(), std::ios::out | std::ios::trunc);
    const size_t data_file_size = FLAGS_simple_recycle_array_test_data_size;
    const size_t doc_num = FLAGS_simple_recycle_array_test_doc_num;
    const size_t init_doc_data_size = data_file_size / doc_num;

    for (size_t idx = 0; idx < doc_num; ++idx) {
        int64_t offset = idx * init_doc_data_size;
        ofs.write((char*)&offset, sizeof(offset));
    }

    ofs.close();
}

int64_t* MmapOffsetFile(const std::string& offset_filename) {
    int32_t fd = open(offset_filename.data(), O_RDWR);
    int64_t* base_addr_ = (int64_t*)mmap(
        nullptr,
        FLAGS_simple_recycle_array_test_doc_num * sizeof(int64_t),
        PROT_READ | PROT_WRITE,
        MAP_SHARED,
        fd,
        0);
    return base_addr_;
}

struct Arg {
    std::vector<std::unique_ptr<hawking::indexlib::SimpleRecycleArray>>* siple_recycle_array = nullptr;
    int64_t* offset_file_addr = nullptr;

    Arg(std::vector<std::unique_ptr<hawking::indexlib::SimpleRecycleArray>>* s, int64_t* o) {
        siple_recycle_array = s;
        offset_file_addr = o;
    }
};

void* Write(void* arg) {
    Arg* sra_arg = reinterpret_cast<Arg*>(arg);
    std::vector<std::unique_ptr<hawking::indexlib::SimpleRecycleArray>>& siple_recycle_array =
        *(sra_arg->siple_recycle_array);
    int64_t* offset_file_addr = sra_arg->offset_file_addr;

    std::uniform_int_distribution<int32_t> dist(0, FLAGS_simple_recycle_array_test_doc_num - 1);
    std::uniform_int_distribution<int32_t> dist2(0, 32);
    std::uniform_int_distribution<int32_t> dist3(0, 16);
    std::random_device rd;
    std::mt19937 mt(rd());
    std::string dummy_data;
    while (1) {
        // clear update data
        dummy_data.clear();

        // random update one doc
        int32_t ran_doc_idx = dist(mt);
        if (nullptr == siple_recycle_array[ran_doc_idx]) {
            siple_recycle_array[ran_doc_idx] = std::make_unique<hawking::indexlib::SimpleRecycleArray>();
        }

        // update-data is a string-list with random count and random-length string
        {
            int32_t str_num = dist2(mt);
            dummy_data.append((const char*)&str_num, sizeof(str_num));
            // LOG(INFO) << "update " << ran_doc_idx << " with " << str_num << " strings";
            for (int32_t idx = 0; idx < str_num; ++idx) {
                int32_t str_len = dist3(mt);
                dummy_data.append((const char*)&str_len, sizeof(str_len));
                // LOG(INFO) << "add str_len " << str_len;
                dummy_data.append(str_len, 'c');
            }
        }

        // upadte to doc's simple_recycle_array
        size_t offset_idx = siple_recycle_array[ran_doc_idx]->Insert(reinterpret_cast<int8_t*>(dummy_data.data()), dummy_data.length());
        while (-1 == offset_idx) {
            offset_idx = siple_recycle_array[ran_doc_idx]->Insert(reinterpret_cast<int8_t*>(dummy_data.data()), dummy_data.length());
            LOG(INFO) << "update doc " << ran_doc_idx << " fail, wait to try again.";
            bthread_usleep(dist2(mt));
        }
        offset_file_addr[ran_doc_idx] = FLAGS_simple_recycle_array_test_data_size + offset_idx;
        
        bthread_usleep(dist2(mt));
    }
}

void* Read(void* arg) {
    Arg* sra_arg = reinterpret_cast<Arg*>(arg);
    std::vector<std::unique_ptr<hawking::indexlib::SimpleRecycleArray>>& siple_recycle_array =
        *(sra_arg->siple_recycle_array);
    int64_t* offset_file_addr = sra_arg->offset_file_addr;

    std::uniform_int_distribution<int32_t> dist(0, 100);
    std::random_device rd;
    std::mt19937 mt(rd());
    while (1) {
        // traverse every doc
        for (int32_t idx = 0; idx < FLAGS_simple_recycle_array_test_doc_num; ++idx) {
            // ignore offset in data-file
            size_t offset = offset_file_addr[idx];
            if (offset < FLAGS_simple_recycle_array_test_data_size) {
                continue;
            }

            // offset in simple_recycle_array, Get the update-data, implicit incr inf
            int8_t* data = siple_recycle_array[idx]->Get(offset - FLAGS_simple_recycle_array_test_data_size);
            EXPECT_NE(nullptr, data);
            {
                // could respect udpate-data with update-bthread
                int32_t* d = reinterpret_cast<int32_t*>(data);
                int32_t str_num = *d;
                // LOG(INFO) << "idx " << idx << " has " << str_num << " strings";
                d += 1;
                for (int32_t idx = 0; idx < str_num; ++idx) {
                    int32_t str_len = *d;
                    d += 1;
                    // LOG(INFO) << "read str_len " << str_len;

                    char* c = reinterpret_cast<char*>(d);
                    for (int32_t len_idx = 0; len_idx < str_len; ++len_idx) {
                        EXPECT_EQ('c', c[len_idx]);
                    }

                    c += str_len;
                    d = reinterpret_cast<int32_t*>(c);
                }
            }

            // release the inf
            siple_recycle_array[idx]->Release(offset - FLAGS_simple_recycle_array_test_data_size);
        }

        bthread_usleep(10);
    }
}

TEST(simple_recycle_array_test, one_write_multi_read) {
    // every doc's SimpleRecycleArray
    std::vector<std::unique_ptr<hawking::indexlib::SimpleRecycleArray>> simple_recycle_array;
    simple_recycle_array.resize(FLAGS_simple_recycle_array_test_doc_num);

    // init offset of every doc
    std::string offset_filename = "/tmp/test_offset_file";
    InitOffsetFile(offset_filename);

    // mmap offset file
    int64_t* offset_file_addr = MmapOffsetFile(offset_filename);
    EXPECT_NE(offset_file_addr, nullptr);

    // bthread arg
    Arg arg(&simple_recycle_array, offset_file_addr);

    // single update thread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, Write, &arg);

    // multi read thread
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_simple_recycle_array_test_tasks);
    for (size_t idx = 0; idx < FLAGS_simple_recycle_array_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, Read, &arg);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }

    munmap(offset_file_addr, FLAGS_simple_recycle_array_test_doc_num * sizeof(int64_t));
}

int main (int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}